#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2021/12/24 15:16
# @Author : xjp
# @Site : 
# @File : generate_dense.py
# @Software: PyCharm
import tensorflow as tf

from ZhengqiLoader import ZhengqiLoader
from generate_CBAM import generate_CBAM


def generate_dense():
    input_xs = tf.keras.Input([5, 5, 1])
    cov_1 = tf.keras.layers.Conv2D(10, 3, padding='SAME', activation=tf.nn.leaky_relu)(input_xs)
    bn1 = tf.keras.layers.BatchNormalization()(cov_1)
    cov_2 = tf.keras.layers.Conv2D(20, 3, padding='SAME', activation=tf.nn.leaky_relu)(bn1)
    bn2 = tf.keras.layers.BatchNormalization()(cov_2)
    cov_3 = tf.keras.layers.Conv2D(30, 3, padding='SAME', activation=tf.nn.leaky_relu)(bn2)
    bn3 = tf.keras.layers.BatchNormalization()(cov_3)
    cov_4 = tf.keras.layers.Conv2D(10, 3, padding='SAME', activation=tf.nn.leaky_relu)(bn3)
    bn4 =tf.keras.layers.BatchNormalization()(cov_4)
    faltten = tf.keras.layers.Flatten()(bn4)
    dense_1 = tf.keras.layers.Dense(40,activation=tf.nn.leaky_relu)(faltten)
    dense_2 = tf.keras.layers.Dense(20,activation=tf.nn.leaky_relu)(dense_1)
    drop_out = tf.keras.layers.Dropout(0.3)(dense_2)
    dense_3 = tf.keras.layers.Dense(10,activation=tf.nn.leaky_relu)(drop_out)
    logits =tf.keras.layers.Dense(6)(dense_3)
    generate = tf.keras.Model(inputs=input_xs, outputs=logits)

    return generate,generate.summary()
if __name__ == '__main__':
    URL = './zhengqi_train.txt'
    BACH_SIZE = 10
    loader = ZhengqiLoader(URL)
    x_train_data, train_dataset, y_train, x_test_data, y_test = loader.preprocess(BACH_SIZE)
    print(x_train_data.shape)
    a, b = generate_dense()

